from core.cuda import cuda_module
from implement.layers.basic.model import Model
from implement.layers.conv.conv2d import Conv2d
from implement.layers.basic.affine import Affine
from utils.functions_collect import relu, flatten, exp


class Encoder(Model):
    def __init__(self, latent_size):
        """编码器模型，用于将输入图像编码为潜在空间中的分布参数。

        包含多个卷积层和仿射变换层，最终输出潜在空间的均值和对数方差。

        Args:
            latent_size (int): 潜在空间的维度大小。
        """
        super().__init__()
        self.latent_size = latent_size
        self.conv1 = Conv2d(32, kernel_size=3, stride=1, pad=1)
        self.conv2 = Conv2d(64, kernel_size=3, stride=2, pad=1)
        self.conv3 = Conv2d(64, kernel_size=3, stride=1, pad=1)
        self.conv4 = Conv2d(64, kernel_size=3, stride=1, pad=1)
        self.linear1 = Affine(32)
        self.linear2 = Affine(latent_size)
        self.linear3 = Affine(latent_size)

    def forward(self, x):
        """编码器的前向传播。

        Args:
            x (ndarray): 输入图像数据。

        Returns:
            tuple: 包含两个元素的元组，分别是潜在空间的均值和对数方差。
        """
        x = relu(self.conv1(x))
        x = relu(self.conv2(x))
        x = relu(self.conv3(x))
        x = relu(self.conv4(x))
        x = flatten(x)
        x = relu(self.linear1(x))
        z_mean = self.linear2(x)
        z_log_var = self.linear3(x)
        return z_mean, z_log_var

    def sampling(self, z_mean, z_log_var):
        """通过潜在空间的均值和对数方差进行采样。

        Args:
            z_mean (ndarray): 潜在空间的均值。
            z_log_var (ndarray): 潜在空间的对数方差。

        Returns:
            ndarray: 采样得到的潜在空间点。
        """
        batch_size = len(z_mean)
        xp = cuda_module
        epsilon = xp.random.randn(batch_size, self.latent_size)
        return z_mean + exp(z_log_var) * epsilon

